# client_eval_remote.py  —— 单步执行版
import time
import base64
import io
import requests
import numpy as np
import torch

from diffusion_policies.common.pytorch_util import dict_apply
from keystroke_counter import KeystrokeCounter, Key, KeyCode
from panda_oyhand_env import PandaOYhandEnv

# ==== 固定超参 ====
DIM_ACTION = 7
N_POINTS = 1024
HORIZON = 8
N_OBS = 2
N_ACTIONS = 5
N_ACTIONS_FINISH = 3
MAX_EPISODE_STEPS = 100
OBS_KEYS = ['point_cloud', 'agent_pos']

# ==== 服务器地址 ====
SERVER_URL = "http://127.0.0.1:18001/predict"

# ==== 编解码 ====
def dumps_npz_b64(arrs):
    buf = io.BytesIO()
    np.savez_compressed(buf, **arrs)
    return base64.b64encode(buf.getvalue()).decode('utf-8')

def loads_npz_b64(b64_str):
    raw = base64.b64decode(b64_str)
    buf = io.BytesIO(raw)
    with np.load(buf, allow_pickle=False) as data:
        return {k: data[k] for k in data.files}

def check_finish(actions):
    print("gripper:", actions[:, 6])
    print("z:", actions[:, 2])
    print([a[6] > 0.03 for a in actions])
    print([a[2] < 0.25 for a in actions])
    stop = np.sum([a[6] > 0.03 and a[2] < 0.25 for a in actions]) >= 1
    if stop:
        print("Task finished!!!")
    # 你原来强制不停止，这里保持一致
    stop = False
    return stop

def request_actions_remote(obs_window):
    payload_b64 = dumps_npz_b64(obs_window)
    resp = requests.post(SERVER_URL, json={'payload_b64': payload_b64}, timeout=10)
    resp.raise_for_status()
    data = resp.json()
    out = loads_npz_b64(data['payload_b64'])
    actions = out['action']  # (N_ACTIONS, DIM_ACTION)
    return actions

def wait_for_step_key(key_counter):
    """阻塞直到按下 空格/回车/n 才继续；按 q 退出（返回 False）"""
    print("↩️  按【空格 / 回车 / n】执行下一步；按【q】退出。", flush=True)
    while True:
        time.sleep(0.03)
        events = key_counter.get_press_events()
        for e in events:
            # 退出
            if e == KeyCode(char='q'):
                return False
            # 继续（空格 / 回车 / 'n'）
            if e in (Key.space, Key.enter) or e == KeyCode(char='n'):
                return True

def main():
    env = PandaOYhandEnv(camera='D435')
    env.go_home()

    try:
        obs = env.reset()
        time.sleep(0.1)
        obs = env.get_obs()

        all_obs_dict = {
            "point_cloud": np.zeros((MAX_EPISODE_STEPS, N_POINTS, 6)),
            "agent_pos": np.zeros((MAX_EPISODE_STEPS, DIM_ACTION))
        }
        all_obs_dict['point_cloud'][0] = obs['point_cloud']
        all_obs_dict['agent_pos'][0] = obs['agent_pos']

        all_actions = np.zeros((MAX_EPISODE_STEPS, DIM_ACTION))
        
        curr_state = obs['agent_pos']

        with KeystrokeCounter() as key_counter:
            stop = False
            finish = False
            action_idx = 1

            while not stop and action_idx < MAX_EPISODE_STEPS:
                print("\n==============================")
                print("action_idx:", action_idx)

                # 随时允许 q 退出
                for e in key_counter.get_press_events():
                    if e == KeyCode(char='q'):
                        stop = True
                        break
                if stop:
                    break

                # 更新观测
                obs = env.get_obs()
                all_obs_dict['point_cloud'][action_idx] = obs['point_cloud']
                all_obs_dict['agent_pos'][action_idx] = obs['agent_pos']

                # 每个窗口起点（1,6,11,...) 从服务器取一块动作
                if action_idx % N_ACTIONS == 1:
                    # 取最近 N_OBS 帧
                    np_obs_dict = {
                        'point_cloud': all_obs_dict['point_cloud'][action_idx - N_OBS + 1:action_idx + 1],
                        'agent_pos':   all_obs_dict['agent_pos'][action_idx - N_OBS + 1:action_idx + 1]
                    }
                    try:
                        action_block = request_actions_remote(np_obs_dict)  # (N_ACTIONS, DIM_ACTION)
                    except Exception as e:
                        print(f"[Client] Remote predict failed: {e}")
                        break
                    
                    action_block_copy = action_block.copy()
                    # for i in range(N_ACTIONS):
                    #     if i == 0:
                    #         action_block_copy[i] = obs['agent_pos'] + action_block_copy[i]
                    #     else:
                    #         action_block_copy[i] = action_block_copy[i] + action_block_copy[i-1]
                    print("action_block_copy:", action_block_copy)
                    all_actions[action_idx:action_idx + N_ACTIONS] = action_block_copy
                    finish = check_finish(action_block_copy)

                    # 如果需要 finish，拼接收尾动作（仍然单步执行）
                    if finish:
                        last_action = all_actions[action_idx + N_ACTIONS - 1].copy()
                        height_offsets = np.zeros((N_ACTIONS_FINISH, DIM_ACTION))
                        height_offsets[:, 2] = np.linspace(0.03, 0.03 * N_ACTIONS_FINISH, N_ACTIONS_FINISH)
                        finish_actions = np.tile(last_action, (N_ACTIONS_FINISH, 1)) + height_offsets
                        finish_actions[:, 5] = -2.4
                        all_actions[action_idx + N_ACTIONS:action_idx + N_ACTIONS + N_ACTIONS_FINISH] = finish_actions

                # print("current_state:", curr_state)
                # print("action_todo:", all_actions[action_idx])
                print('delta_action:', np.round(all_actions[action_idx] - curr_state, 3))
                # === 单步等待：按键后才执行当前一步 ===
                proceed = wait_for_step_key(key_counter)
                if not proceed:
                    print("收到 q，退出。")
                    break

                # action_todo: 必须是绝对位置
                action_todo = all_actions[action_idx]
                
                if action_todo[2] < 0.11:
                    action_todo[2] = 0.11
                obs, reward, done, info = env.step(action_todo)
                
                curr_state = action_todo.copy()

                action_idx += 1

    finally:
        env.stop()

if __name__ == "__main__":
    main()